#!/bin/bash
#Author:IGI
#Email:chenwqin@gmail.com
#Date: 2009-10-23
#Modify:2010-10-08
#License:GPLv3

#myssh version
version="1.5.0"

#scriptname, 非必要变量
scriptname="$0"

#format: class1:class2:...:IP#Sname#User#Port#Description
#ip列表文件, 支持多列表, 以空格分隔
servers_list="./serverlist"

#日志记录位置, 此变量必须是目录
log="/tmp"
error_iplist="/tmp/myssh.error.`date +%Y%m%d%H%M%S`"
lock="/tmp/myssh.pid"

#日志记录方法, 0为每次运行程序清除上一次log, 1为不清除, 非必要变量
safelog=1

default_user=`whoami`
default_port='22' #非必要变量

#terminal支持, 可以使用以下值， multixterm没有经过详细测试
#konsole, gnome-terminal, xterm, multixterm, cssh, 'konsole cssh'
default_terminal='konsole cssh'
default_description='' #非必要变量

#警告tab数及最大tab数
warning_tab_count=50
max_tab_count=200

#串行运行时的进程数
pnum=50

#修改默认执行方式, 1采用并行
ifp=1

#调试开关, 非必要变量
debug=1

#ssh命令格式
ssh_command='ssh -oStrictHostKeyChecking=no $ssh_options -p $port $user@$ip '

#screen命令
screen_command='screen -S $ip' 
    
#变量检查函数
check_var(){
for i in $@
do
    if [ -z "${!i}" ]
    then
        echo "var check:$i error"
        exit 99
    fi
done
}

#检查是否有其他myssh在运行
check_lock(){
pid=`<$lock`
if [ -n "$pid" ]
then
    ps -p "$pid" | grep myssh
    if [ $? -eq 0 ]
    then
        echo 'Another myssh still running'
        exit
    fi
fi
echo $$ > $lock
}

#检查错误的服务器IP
check_error_iplist(){
    if [ -f "$error_iplist" ]
    then
        errorip=`<$error_iplist`
        if [ -n "$errorip" ]
        then
            echo "Merge error ip lists..."
            echo -e "$errorip" | sed -e '/^$/d' -e "s/$/\terror/"
        fi
    fi
}

#检查命令执行状态
check_exit(){
    if [ $? -ne 0 ] 
    then
        if [ ${debug:-0} -eq 1 ]
        then
            echo -en "$ip\t""`<$log/$ip`""\n"
        else
            echo -e "$ip\terror"
        fi
        echo "$ip" >> "$error_iplist" 
    else
        echo -en "$ip\t""`<$log/$ip`""\n"
    fi    
}


#判断用户是否决定
check_ans(){
    read -p "$@" ans
    case $ans in 
    n|N|no|No|NO)
        echo "Exit done"
        exit 10
        ;;
    *)
        return 
        ;;
    esac
}

#判断是否开启safe_log模式
check_log(){
    if [ ${safelog:-0} -eq 1 ]
    then
        log="$log-`date +%Y%m%d-%H-%M-%S`"
        mkdir "$log"
    else
        rm -rf $log
        mkdir "$log"
    fi
}


#检查用户当前是否是终端登录和远程登录
check_tty(){
    tmp_info=`who -m | awk '{print $5}'`
    if [[ ! "$tmp_info" =~ "(:" ]]
    then
        specify_new_windows=0
    fi
}


#判断ip列表中ip是否有重复且信息不同, 以避免同一服务器多次执行命令的危险
check_ip_same(){
    allcount=`echo "$1" | sort -u | wc -l`
    ipcount=`echo "$1" | awk -F# '{print $1}' | sort -u | wc -l`
    if [ $allcount -ne $ipcount ]
    then
        echo "Iplist conflict!"
        echo "$1"
        exit 33
    fi
}

#判断参数个数
check_option_number(){
    #$1 option name
    #$2 option number
    #$3 option expect number
    if [ $2 -lt $(($3+1)) ]
    then
        echo "$scriptname:option $1 requires $3 arguments at least"
        exit 0
    fi
}
    

#检查ip列表文件是否存在
check_list(){
for i in $servers_list
do
    if [ ! -r "$i" ]
    then
        if [ -e "$i" ]
        then
            echo "Can not open list file:$i"
            exit 2
        else
            echo "Can not find list file:$i"
            exit 3
        fi
    fi
done
list=`cat $servers_list | grep -v '^#\|^$'`
#临时切割用变量
reslist="$list"
}


#检查terminal变量
check_terminal(){
    if ! which $default_terminal >/dev/null 2>&1
    then
        echo "$default_terminal isn't exist."
        exit 4
    fi
        
    if [ "$default_terminal" == "konsole" ]
    then
        command_string="konsole --new-tab --hold -e $screen_command $@" 
    elif [ "$default_terminal" == "gnome-terminal" ]
    then
        command_string="gnome-terminal --tab-with-profile='linux' -e \""$screen_command $@"\""
    elif [ "$default_terminal" == "xterm" ]
    then
        command_string="xterm -e \""$screen_command $@"\"" 
    elif [ "$default_terminal" == "multixterm" ]
    then
        command_string="multixterm -xc \""%n"\" \""$@"\""

    elif [ "$default_terminal" == "konsole cssh" ]
    then
        command_string="konsole --new-tab --hold -e $screen_command $@"
        command_string_cssh="cssh" 
    #Modify here when your favorite terminal isn't gnome-terminal or kde konsole
    #elif [ "$default_terminal" == "your favorite terminal" ]
    #then
    #    command_string="your favorite terminal options"
    else
        echo "Please modify the check_terminal function in myssh"
        exit 5 
    fi
}


#检查将要打开的tab数量
check_tabcount(){
    count=`echo "$1" | awk -F: '{print $NF}' | awk -F# '{print $1}' | sort -u | wc -l`
    if [ $count -gt $max_tab_count ]
    then
        echo "You want to open $count tabs and it is not allowed"
        exit 13
    elif [ $count -gt $warning_tab_count ]
    then
        check_ans "You want to open $count tabs,are you sure?(y/n):"
    fi
    return $count
}

#设置IP变量, port变量, user变量
set_var(){
    #不用数组是因为在用户名为空时会错误的把端口当成用户名
    ip=`echo "$@" | awk -F# '{print $1}'`
    list_user=`echo "$@" | awk -F# '{print $3}'`
    list_port=`echo "$@" | awk -F# '{print $4}'`
    user="${req_user:-${list_user:-$default_user}}"
    port="${req_port:-${list_port:-$default_port}}"
}

#为ip排序
sort_iplist(){
    input=`cat - | sort -u`
    ifip=`echo "$input" | grep -c '^[0-9]\{1,3\}\.[0-9]\{1,3\}\.[0-9]\{1,3\}\.[0-9]\{1,3\}'`
    if [ "$ifip" -eq 0 ]
    then
        echo "$input"
    else
        echo "$input" | sort -t. -k1,1n -k2,2n -k3,3n -k4,4n
    fi
}

#切割生成需要显示的select选项 
display(){
    echo "$reslist" | cut -d: -f1 | awk -F# '{if($2 != ""){print $1"#"$2}else{print $1}}' | sort_iplist
}

#cut iplist
cut_list(){
    echo "$1" | awk -v GN="$2" -F: '$1==GN{print $0}' | cut -d: -f2-
}

#menu
menu(){
    select i
    do
        if [ -z "$i" ]
        then
            echo "Empty select!"
            continue
        else
            reslist=`cut_list "$reslist" "$i"`
            if [ -z "$select_res" ]
            then
                select_res="$i"
            else
                select_res="$select_res:$i"
            fi
            menu $(display) 
            return
        fi  
    done
}

#grep iplist from serverlist
grep_iplist(){
    greptmp=`echo "$list" | cut -d# -f1-4`
    for i in $@
    do
        #加斜杆避免grep出错
        greptmp=`echo "$greptmp" | grep -w "\$i"`
    done
    echo "$greptmp" | awk -F: '{print $NF}'
}

#不严格匹配字段位置, 显示全信息
search_iplist(){
    cut_res=0
    if [ "$1" == '-ip' ] || [ "$1" == '--ipaddress' ]
    then
        shift
        cut_res=1
    fi

    searchtmp="$list"
    for i in $@
    do
        #加斜杆避免grep出错, 不使用grep可以提高效率
        #searchtmp=`echo "$searchtmp" | fgrep -w "\$i"`
        searchtmp=`echo "$searchtmp" | sed -n "/\<$i\>/p"`
    done

    if [ -z "$searchtmp" ]
    then
        echo "Empty search!"
        exit 11
    fi
    iplist=`echo "$searchtmp" | awk -F: '{split($NF,last,"#");print last[1]}'`
    #如果-ip参数,　则只显示IP地址
    if [ "$cut_res" -eq 1 ]
    then
        echo "$iplist" | sort_iplist
    else
        searchtmp=`echo "$searchtmp" | sed 's/:\|#/ /g' | sort -u | column -t`
        echo "$searchtmp" 
        echo "Total count $(echo "$searchtmp" | wc -l)"
        echo "Total ipadd $(echo "$iplist" | wc -l)"
    fi
    exit 0
}


#expect执行方式, 由su变量及ifssh变量同时决定是否启用此方式
expect_run_command(){
expect <<EOF
set timeout 30
log_user 0 
spawn $ssh_command_res
expect {
       "*\$*" {send "\n"}
       "*(yes/no)?" {send "yes\n"}
       "*Permission denied*" {exit}
}
expect "*\$*" {send "su\n"}
expect "*Password*" {send "$password\n"}
expect {
        "*\$*" {
                send "su\n"
                expect "*Password*" {send "$password\n"
                                     expect "*\$*" {exit}
                                    }
        }
        "*#*" {send "\n"}
}
expect "*#*" {send "$ssh_exec >$log 2>&1\n"}
expect "*#*" {send "exit\n"}
expect "*\$*" {send "exit\n"}
EOF
}

#执行函数, ifp变量决定是否采用并行模式, 默认并行
#new_windows变量决定是否新开窗口
run_command(){
    if [ -z "$1" ] 
    then
        echo "Select empty!"
        exit 0
    fi
    check_ip_same "$1"
    check_tty

    #并行可新开tab函数
    if [ ${specify_new_windows:-${new_windows:-0}} -eq 1 ]
    then
        #echo 'mode pa=? nw=1'
        #exit 1
        check_tabcount "$1"
        count=$?
        
        run_list=`echo "$1" | sort -u`
        shift
        check_terminal $@
        
        if [ ${_r:-0} -eq 0 ]
        then
            if [ $count -gt 1 ] && [ $count -le 33 ] && [[ "$default_terminal" =~ 'cssh' ]] 
            then
                for i in $run_list
                do
                    set_var "$i"
                    ip_res="$ip_res $user@$ip:$port"
                done
                eval "$command_string_cssh $ip_res" >/dev/null &
                #eval "$command_string_cssh $ip_res" 
            else
                for i in $run_list
                do
                    #设置生成ip,port,user
                    set_var "$i"
                    echo "$ip"
                    eval "$command_string" &
                    #eval echo "$command_string"
                    if [ "$?" -ne 0 ] && [ ${debug:-0} -eq 1 ]
                    then
                        eval echo "$1:$command_string"
                    fi
                    #避免konsole新开tab过快而出错
                    sleep 0.1
                done
            fi
        else
            command_string=`echo $command_string | sed 's/screen -S $ip//g'`
            for id in $run_list
            do
                echo $id
                $command_string $id &
                if [ "$?" -ne 0 ] && [ ${debug:-0} -eq 1 ]
                then
                    eval echo "$1:$command_string $id"
                fi
            done
        fi
    else
        #并行,不新开tab函数
        if [ ${specify_ifp:-$ifp} -eq 1 ]
        then
            #echo 'mode pa=1 nw=0'
            #exit 1
            run_list=`echo "$1" | sort -u`
            shift
            command_string="${@}"
            #echo "$command_string"
           
            command_string=`echo "$command_string"` 
            ssh_command=`echo "$ssh_command"`

            #多进程数量控制
            tmpfife="$$.fifo"
            mkfifo /tmp/$tmpfife
            exec 4<>/tmp/$tmpfife
            rm -f /tmp/$tmpfife
            for ((i=0;i<$pnum;i++))
            do
                echo >&4
            done
    
            #检查是否启用safe_log
            check_log
            
            for i in $run_list
            do
                read
                set_var "$i"
                
                #是否切换到root权限 
                if [ ${su:-0} -eq 1 ] && [ ${ifssh:-0} -eq 1 ]
                then
                    ssh_command_res=`eval echo "$ssh_command"`
                    { expect_run_command > $log/$ip 2>&1 && $ssh_command_res "cat $log"  >> $log/$ip 2>&1 ;
                      check_exit;
                      sleep 0.`date +%N` && echo >&4 ;} < /dev/null &
                else
                    #sleep解决占位符丢失问题
                    #< /dev/null 解决read报错问题
                    #eval echo "$command_string"
                    { eval "$command_string" > $log/$ip 2>&1 ;
                      check_exit;
                      sleep 0.`date +%N` && echo >&4 ;} < /dev/null &
                fi

            done <&4
            wait
            exec 4>&-
        else
            #完全串行函数, 不新开tab
            #echo 'mode pa=0 nw=0'
            #exit 1
            run_list=`echo "$1" | sort -u`
            shift
            command_string=${@}

            #检查是否启用safe_log
            check_log

            command_string=`echo "$command_string"` 
            ssh_command=`echo "$ssh_command"`
        
            for i in $run_list
            do
                set_var "$i"

                #判断是否启用root模式及是否是执行ssh命令
                if [ ${su:-0} -eq 1 ] && [ ${ifssh:-0} -eq 1 ]
                then
                    ssh_command_res=`eval echo "$ssh_command"`
                    expect_run_command > $log/$ip 2>&1 && $ssh_command_res "cat $log"  > $log/$ip 2>&1
                else
                    eval "$command_string" | tee -a $log/$ip 
                    #eval echo "$command_string"
                fi
                check_exit
            done
        fi
    fi
    check_error_iplist
}

#scp参数处理函数
scp_command(){
    echo "$@" | sed 's/^/scp -q -P $port /' | sed 's/:\(.*\)/$user@$ip:\1 \&\& echo done || echo error/'
}

#给出列表选择菜单, 并根据选择的机器执行ssh
empty_option(){
    menu $(display)
    
    if [ -z "$select_res" ]
    then
        exit 0
    fi
    
    echo "$select_res" 
    check_ans "Yes or no?:"
    res=`grep_iplist "^$select_res"`
    new_windows=1
    ifp=0
    run_command "$res" 'ssh -l $user -p $port $ip'
}


#显示帮助手册
help_menu(){
    echo "Usage: myssh [[group] [class] [sname] [ip]]"
    echo "             [-l|--list [-ip|--ipaddress] [-r] [string]]"
    echo "             [-ip|--ipaddress [string]]"
    echo "             [-e|--exec 'command \$ip|\$user|\$port' string]"
    echo "             [-re|--remote-exec 'command' string]"
    echo "             [-are|--ask-remote-exec string]"
    echo "             [-s|--scp 'options file :path'|'options :path file' string]"
    echo "             [-p|--ping string]"
    echo "             [-kp|--quick-ping string]"
    echo "             [-t|--traceroute string]"
    echo "             [-mtr|--mtr string]"
    echo "             [-P|--port port]"
    echo "             [-A]"  #ssh -A 
    echo "             [-D [bind_address:]port]" #ssh -D
    echo "             [-np|--no-parallel]" #turn off parallel mode
    echo "             [-pnum|--parallel-numbers numbers]"
    echo "             [-root|--root [password] -re 'command' string]" #remote-exec with su to root
    echo "             [-sl|--safelog]" #open safelog mode
    echo "             [-u|--user user]"
    echo "             [-m|--modify]"
    echo "             [-i|--input iplist [options]]"
    echo "             [-o|--option var=value]"
    echo "             [-r|--resumes]" #list or resumes  all  detached  screen session
    echo "             [-h|--help]"
    echo "             [-v|--version]"
    exit 0
}


#自动补全输出
#需要在bashrc文件中自行添加: complete -W "`myssh -c`" -o default myssh
completion(){
   completion_options='-l --list -ip --ipaddress -e --exec -re --remote-exec -are --ask-remote-exec -root --root -s --scp -p --ping -kp --quick-ping -t --traceroute -mtr --mtr -P --port -A -D -np --no-parallel -pnum --parallel-numbers -sl --safelog -u --user -m --modify -i --input -h --help -v --version -o --option -r --resumes' 
   completion_var='new_windows= ifp= log= safelog= user= port= default_terminal= warning_tab_count= max_tab_count= pnum= debug= ssh_command= ssh_options= req_port= req_user= ifssh= error_iplist=' 
   completion_iplist=`echo "$list" | awk -F[:#] '{for(i=1;i<NF-2;i++)print $i}' | sort -u`  
   echo "$completion_iplist" "$completion_options" "$completion_var"
}

#根据参数执行
run_options(){
xnum="${_h:-0}${_v:-0}${_l:-0}${_ip:-0}${_e:-0}${_re:-0}${_are:-0}${_p:-0}${_kp:-0}${_t:-0}${_mtr:-0}${_s:-0}${_m:-0}${_c:-0}${_r:-0}"
#010101010101010
case $xnum in 
    100000000000000)
        help_menu
        ;;
    010000000000000)
        echo "Myssh version v$version"
        ;;
    110000000000000)
        echo "Myssh version v$version"
        help_menu
        ;;
    001000000000000)
        search_iplist $res
        ;;
    001100000000000)
        search_iplist '-ip' $res
        ;;
    000100000000000)
        search_iplist '-ip' $res
        ;;
    000010000000000)
        res=`grep_iplist $res`
        run_command "$res" "$to_exec"
        ;;
    000001000000000)
        res=`grep_iplist $res`
        if echo "$ssh_exec" | grep -q "'"
        then
            run_command "$res" "$ssh_command \"$ssh_exec\""
        else
            run_command "$res" "$ssh_command '$ssh_exec'"
        fi    
        ;;
    000000100000000)
        echo 'Input your command(End of ";")'
        read -d';' ssh_exec
        echo 
        res=`grep_iplist $res`
        if echo "$ssh_exec" | grep -q "'"
        then
            run_command "$res" "$ssh_command <<<\"$ssh_exec\""
        else
            run_command "$res" "$ssh_command <<<\"$ssh_exec\""
        fi    
        ;;
    000000010000000)
        res=`grep_iplist $res`
        run_command "$res" 'ping $ip'
        ;;
    000000001000000)
        res=`grep_iplist $res`
        run_command "$res" '{ ping -A -c1 $ip || ping -c2 $ip ; } | grep -i from'
        if [ -n "$errorip" ]
        then
            echo "Searching ip lists..."
            for i in `echo -e "$errorip"`
            do
                $0 -l "$i" | grep -v 'Total'
            done  | sort -u | column -t
        fi
        ;;
    000000000100000)
        res=`grep_iplist $res`
        run_command "$res" 'traceroute $ip'
        ;;
    000000000010000)
        res=`grep_iplist $res`
        run_command "$res" 'sudo mtr -t $ip'
        ;;
    000000000001000)
        res=`grep_iplist $res`
        run_command "$res" "$to_exec"
        ;;
    000000000000100)
        select i in $servers_list
        do
            if [ -z "$i" ]
            then
                echo "Empty select!"
                continue
            else
                vim $i
                exit
            fi  
        done
        ;;
    000000000000010)
        completion
        ;;
    000000000000001)
        screen_list=`screen -ls | grep -i detached`
        for i in $res
        do
            screen_list=`echo "$screen_list" | grep -w $i`
        done
        res=`echo "$screen_list" | awk -F. '{print $1}' | sed 's/\s//g'`
        new_windows=1
        ifp=0
        run_command "$res" 'screen -r'
        ;;
    001000000000001)
        screen -ls | grep -i detached
        ;;
    000000000000000)
        res=`grep_iplist $res`
        new_windows=1
        ifp=0
        run_command "$res" "$ssh_command"
        ;;

    *)
        echo "Error Usage"
        exit 1
        ;;
esac
}


#参数检测函数
check_options(){
if [ -n "$1" ]
then
    case "$1" in
        -*)
            case "$1" in
                -h|--help) 
                    _h=1
                    shift
                    check_options "$@"
                    ;;
                -v|--version)
                    _v=1
                    shift   
                    check_options "$@"
                    ;;
                -r|--resumes)
                    #list or resumes  all  detached  screen session
                    _r=1
                    shift
                    check_options "$@"
                    ;;
                -l|--list)
                    _l=1
                    shift
                    check_options "$@"
                    ;;
                -ip|--ipaddress)
                    _ip=1
                    shift
                    check_options "$@"
                    ;;
                -e|--exec)
                    _e=1
                    shift
                    to_exec="$1"
                    shift
                    check_options "$@"
                    ;;
                -re|--remote-exec)
                    _re=1
                    shift
                    ifssh=1
                    ssh_exec="$1"
                    shift
                    check_options "$@"
                    ;;

                -are|--ask-remote-exec)
                    _are=1
                    shift
                    ifssh=1
                    check_options "$@"
                    ;;
                -root|--root)
                    su=1
                    read -p "Password:" -s password
                    shift
                    check_options "$@"
                    ;;
                -p|--ping)
                    _p=1
                    new_windows=1
                    ifp=0
                    shift
                    check_options "$@"
                    ;; 
                -kp|--quick-ping)
                    _kp=1
                    shift
                    check_options "$@"
                    ;;
                -t|--traceroute|--tracert)
                    _t=1
                    new_windows=1
                    ifp=0
                    shift
                    check_options "$@"
                    ;;
                -mtr|--mtr)
                    _mtr=1
                    new_windows=1
                    ifp=0
                    shift
                    check_options "$@"
                    ;;
                -s|--scp)
                    _s=1
                    shift
                    to_exec=`scp_command "$1"`
                    shift
                    check_options "$@"
                    ;;
                -o|--option)
                    shift
                    eval $1
                    shift
                    check_options "$@"
                    ;;
                -P|--port)
                    shift
                    req_port="$1"
                    shift
                    check_options "$@"
                    ;;
                -u|--user)
                    shift
                    req_user="$1"
                    shift
                    check_options "$@"
                    ;;
                -A)
                    ssh_options="$ssh_options $1"
                    shift
                    #ssh-agent
                    eval `ssh-agent`
                    ssh-add
                    check_options "$@"
                    ;;
                -D)
                    ssh_options="$ssh_options $1"
                    shift
                    ssh_options="$ssh_options $1"
                    shift
                    check_options "$@"
                    ;;
                -np|--no-parallel)
                    #exec command with no parallel mode
                    specify_ifp=0
                    specify_new_windows=0
                    shift
                    check_options "$@"
                    ;;
                -pnum|--parallel-numbers)
                    shift
                    pnum=$1
                    shift
                    check_options "$@"
                    ;; 
                -sl|--safelog)
                    safelog=1
                    shift
                    check_options "$@"
                    ;;
                -m|--modify)
                    _m=1
                    shift
                    check_options "$@"
                    ;;
                -i|--input)
                    shift
                    if [ -f "$1" ]
                    then
                        list=`<$1`
                    else
                        list=`echo "$1"`
                    fi

                    reslist="$list"
                    shift
                    check_options "$@" '[0-9]\{1,3\}\.[0-9]\{1,3\}\.[0-9]\{1,3\}\.[0-9]\{1,3\}'
                    ;;
                -c)
                    _c=1
                    shift
                    check_options "$@"
                    ;;
                *)
                    echo "Unknown Option '$1'"
                    help_menu
                    exit 9
                    ;;
            esac
            ;;
        *)
            res="$res $1"
            shift
            check_options "$@"
            ;;
    esac
fi
}

#主函数
main(){
    check_var servers_list log default_terminal warning_tab_count max_tab_count pnum ifp ssh_command error_iplist
    check_list
    screen -wipe >/dev/null 2>&1
    log="$log/myssh-log"

    if [ "$#" -eq 0 ]
    then
        #参数为空, 执行菜单选择函数
        empty_option
        exit 0
    fi
    
    check_options "$@"
    run_options
}

main "$@"
